#include "unetv1.h"
#include "stream_tools.h"
#include <vector>

template<
    unsigned BATCH,
    unsigned IN_HEIGHT,
    unsigned IN_WIDTH,
    unsigned IN_CHANNEL,
    unsigned KERNEL,
    unsigned STRIDE,
    unsigned PAD,
    unsigned DWIDTH,
    unsigned P_CH,
    unsigned PADVAL
>
void im2col_golden(data_stream<DWIDTH*P_CH>& in, data_stream<DWIDTH*P_CH>& out)
{
    static_assert(IN_CHANNEL >= P_CH, "im2col_golden");
    static_assert(IN_CHANNEL % P_CH == 0, "im2col_golden");
    constexpr unsigned FOLD = IN_CHANNEL / P_CH;
    constexpr unsigned OUT_HEIGHT = (IN_HEIGHT + 2 * PAD - KERNEL) / STRIDE + 1;
    constexpr unsigned OUT_WIDTH = (IN_WIDTH + 2 * PAD - KERNEL) / STRIDE + 1;
    const ap_uint<DWIDTH*P_CH> PAD_AP = 0;
    for(unsigned i = 0; i < P_CH; ++i) {
        PAD_AP.range((i+1)*DWIDTH-1, i*DWIDTH) = PADVAL;
    }

    std::cout << "im2col_golden: " 
              << "BATCH: " << BATCH 
              << ", IN_HEIGHT: " << IN_HEIGHT 
              << ", IN_WIDTH: " << IN_WIDTH 
              << ", KERNEL: " << KERNEL 
              << ", STRIDE: " << STRIDE 
              << ", PAD: " << PAD 
              << ", FOLD: " << FOLD 
              << std::endl;
    
    assert(in.size() == BATCH * IN_HEIGHT * IN_WIDTH * FOLD);
    assert(out.empty());
    
    // Read all input data into a buffer (use dynamic allocation to avoid stack overflow)
    std::vector<std::vector<std::vector<std::vector<ap_uint<DWIDTH*P_CH>>>>> input_buf(
        BATCH, std::vector<std::vector<std::vector<ap_uint<DWIDTH*P_CH>>>>(
            IN_HEIGHT, std::vector<std::vector<ap_uint<DWIDTH*P_CH>>>(
                IN_WIDTH, std::vector<ap_uint<DWIDTH*P_CH>>(FOLD)
            )
        )
    );
    
    for (unsigned b = 0; b < BATCH; ++b) {
        for (unsigned h = 0; h < IN_HEIGHT; ++h) {
            for (unsigned w = 0; w < IN_WIDTH; ++w) {
                for (unsigned f = 0; f < FOLD; ++f) {
                    input_buf[b][h][w][f] = in.read();
                }
            }
        }
    }

    // Generate im2col output
    for (unsigned b = 0; b < BATCH; ++b) {
        for (unsigned oh = 0; oh < OUT_HEIGHT; ++oh) {
            for (unsigned ow = 0; ow < OUT_WIDTH; ++ow) {
                for (unsigned f = 0; f < FOLD; ++f) {
                    for (unsigned kh = 0; kh < KERNEL; ++kh) {
                        for (unsigned kw = 0; kw < KERNEL; ++kw) {
                            int ih = oh * STRIDE + kh - PAD;
                            int iw = ow * STRIDE + kw - PAD;
                            
                            if (ih < 0 || ih >= IN_HEIGHT || iw < 0 || iw >= IN_WIDTH) {
                                out.write(PAD_AP);
                            } else {
                                out.write(input_buf[b][ih][iw][f]);
                            }
                        }
                    }
                }
            }
        }
    }
    
    assert(in.empty());
    assert(out.size() == BATCH * OUT_HEIGHT * OUT_WIDTH * FOLD * KERNEL * KERNEL);
    return;
}

template<
    unsigned BATCH,
    unsigned IN_HEIGHT,
    unsigned IN_WIDTH,
    unsigned IN_CHANNEL,
    unsigned KERNEL,
    unsigned STRIDE,
    unsigned PAD,
    unsigned DWIDTH, // Data width in bits
    unsigned P_CH,
    unsigned PADVAL
>
void im2col(data_stream<DWIDTH*P_CH>& in, data_stream<DWIDTH*P_CH>& out)
{
    static_assert(IN_CHANNEL >= P_CH, "im2col");
    static_assert(IN_CHANNEL % P_CH == 0, "im2col");
    constexpr unsigned FOLD = IN_CHANNEL / P_CH;
    constexpr unsigned OUT_HEIGHT = (IN_HEIGHT + 2 * PAD - KERNEL) / STRIDE + 1;
    constexpr unsigned OUT_WIDTH = (IN_WIDTH + 2 * PAD - KERNEL) / STRIDE + 1;
    constexpr unsigned PRELOAD_PIXELS = KERNEL - 1;
    const ap_uint<DWIDTH*P_CH> PAD_AP = 0;
    for(unsigned i = 0; i < P_CH; ++i) {
        PAD_AP.range((i+1)*DWIDTH-1, i*DWIDTH) = PADVAL;
    }
    std::cout << "im2col: " 
              << "BATCH: " << BATCH 
              << ", IN_HEIGHT: " << IN_HEIGHT 
              << ", IN_WIDTH: " << IN_WIDTH 
              << ", KERNEL: " << KERNEL 
              << ", STRIDE: " << STRIDE 
              << ", PAD: " << PAD 
              << ", FOLD: " << FOLD 
              << std::endl;
    assert(in.size() == BATCH * IN_HEIGHT * IN_WIDTH * FOLD);
    assert(out.empty());

    #pragma HLS DATAFLOW

    ap_uint<DWIDTH*P_CH> line[KERNEL][IN_WIDTH][FOLD];
    #pragma HLS ARRAY_PARTITION variable=line complete dim=1

    for (unsigned iw_idx = 0; iw_idx < IN_WIDTH; ++iw_idx) {
        for (unsigned f = 0; f < FOLD; ++f) {
        #pragma HLS PIPELINE II=1
            line[0][iw_idx][f] = in.read();
        }
    }

    for (unsigned iw_idx = 0; iw_idx < PRELOAD_PIXELS; ++iw_idx) {
        for (unsigned f = 0; f < FOLD; ++f) {
        #pragma HLS PIPELINE II=1
            line[1][iw_idx][f] = in.read();
        }
    }

    // main loop
    for (unsigned b_idx= 0; b_idx< BATCH; ++b_idx) {
        for (unsigned ih_idx = 0; ih_idx < IN_HEIGHT; ++ih_idx) {
            for (unsigned iw_idx = 0; iw_idx < IN_WIDTH; ++iw_idx) {
                for (unsigned f = 0; f < FOLD; ++f) {
                    #pragma HLS PIPELINE II=1
                    unsigned ih_idx_to_write = ih_idx + 1;
                    unsigned iw_idx_to_write = iw_idx + PRELOAD_PIXELS;
                    if(iw_idx_to_write >= IN_WIDTH) {
                        iw_idx_to_write = iw_idx + PRELOAD_PIXELS - IN_WIDTH;
                        ih_idx_to_write += 1;
                    }
                    if(ih_idx_to_write < IN_HEIGHT) {
                        ap_uint<DWIDTH*P_CH> in_buf = in.read();
                        unsigned line_idx_to_write = ih_idx_to_write % 3;
                        line[line_idx_to_write][iw_idx_to_write][f] = in_buf;
                    }
                    if(ih_idx % STRIDE == 0 && iw_idx % STRIDE == 0) {
                        unsigned oh_idx = ih_idx / STRIDE;
                        unsigned ow_idx = iw_idx / STRIDE;
                        // 因为一开始预加载了若干行和额外的几个pixel，所以主循环中加载时的坐标需要调整
                        for(unsigned kh = 0; kh < KERNEL; ++kh) {
                            for(unsigned kw = 0; kw < KERNEL; ++kw) {
                                unsigned ih_idx_to_read = oh_idx * STRIDE + kh - PAD;
                                unsigned iw_idx_to_read = ow_idx * STRIDE + kw - PAD;
                                if (ih_idx_to_read < 0 || ih_idx_to_read >= IN_HEIGHT ||
                                    iw_idx_to_read < 0 || iw_idx_to_read >= IN_WIDTH) {
                                    out.write(PAD_AP);
                                } else {
                                    unsigned line_idx_to_read = ih_idx_to_read % 3;
                                    out.write(line[line_idx_to_read][iw_idx_to_read][f]);
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    assert(in.empty());
    assert(out.size() == BATCH * OUT_HEIGHT * OUT_WIDTH * FOLD * KERNEL * KERNEL);
    return;
}
